use bytes::{Buf, BufMut, Bytes, BytesMut};

use crate::prelude::{Decoder, Encoder, MessageId};

pub struct UserRawMsg {
    pub id: u16,
    pub bytes: Bytes,
}

pub struct UserMsgCodec {}

impl UserMsgCodec {
    pub fn new() -> Self {
        Self {}
    }
}

impl<'a, T> Encoder<&'a T> for UserMsgCodec
where
    T: 'static + MessageId + prost::Message,
{
    type Error = anyhow::Error;

    fn encode(&mut self, item: &'a T, dst: &mut BytesMut) -> Result<(), Self::Error> {
        let length = super::U16_LENGTH + super::U16_LENGTH + item.encoded_len();
        dst.reserve(length);
        dst.put_u16((length - super::U16_LENGTH) as _);
        dst.put_u16(T::ID);
        item.encode(dst)?;
        Ok(())
    }
}

impl Encoder<UserRawMsg> for UserMsgCodec {
    type Error = anyhow::Error;

    fn encode(&mut self, item: UserRawMsg, dst: &mut BytesMut) -> Result<(), Self::Error> {
        let length = super::U16_LENGTH + super::U16_LENGTH + item.bytes.len();
        dst.reserve(length);
        dst.put_u16((length - super::U16_LENGTH) as _);
        dst.put_u16(item.id);
        dst.put(item.bytes);
        Ok(())
    }
}

impl Decoder for UserMsgCodec {
    type Item = UserRawMsg;
    type Error = anyhow::Error;

    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        if src.len() < super::U16_LENGTH {
            return Ok(None);
        }
        let mut len_bytes = [0u8; super::U16_LENGTH];
        len_bytes.copy_from_slice(&src[..super::U16_LENGTH]);
        let len = u16::from_be_bytes(len_bytes) as usize;

        if src.len() - super::U16_LENGTH >= len {
            // Skip the length header we already read.
            src.advance(super::U16_LENGTH);
            let mut bytes = src.split_to(len).freeze();
            let id = bytes.get_u16();
            Ok(Some(UserRawMsg { id, bytes }))
        } else {
            Ok(None)
        }
    }
}
